# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

import os
import torch
import numpy as np
import pandas as pd

from detectron2.data import DatasetCatalog, MetadataCatalog
from tqdm import tqdm


__all__ = ["register_ava"]


# fmt: off
CLASS_NAMES = [
    "photo"
]
# fmt: on


def load_ava_instances(dirname: str, split: str):
    """
    Load Pascal ava detection annotations to Detectron2 format.

    Args:
        dirname: Contain "Annotations", "ImageSets", "JPEGImages"
        split (str): one of "train", "test", "val", "trainval"
    """

    if split == "val":
        image_base = os.path.join(dirname, "val")
        avail_files = os.listdir(image_base)
        dicts = []
        for idx, im in enumerate(avail_files):
            r = {
                "file_name": os.path.join(image_base, im),
                "image_id": idx,
                "height": 1,
                "width": 1,
            }
            r["label"] = torch.ones((10,))
            dicts.append(r)
        return dicts

    cache_dir = os.path.join(dirname, f"ava_{split}.pkl")
    if os.path.exists(cache_dir):
        with open(cache_dir, 'rb') as f:
            return torch.load(f)

    with open(os.path.join(dirname, 'AVA.txt')) as f:
        df = torch.tensor([[int(i) for i in l.split()] for l in f.readlines()])
    image_base = os.path.join(dirname, "images")
    set_map = {
        "train": "generic_ls_train.jpgl",
        "test": "generic_test.jpgl",
    }

    avail_files = os.listdir(image_base)
    dicts = []
    if split == "full":
        for d in tqdm(df):
            ImageID = d[1]
            jpeg_file = os.path.join(image_base, f"{ImageID}.jpg")
            if f"{ImageID}.jpg" in avail_files:
                r = {
                    "file_name": jpeg_file,
                    "image_id": ImageID,
                    "height": 1,
                    "width": 1,
                }
                r["label"] = d[2:12]
                dicts.append(r)
    else:
        with open(os.path.join(dirname, "aesthetics_image_lists", set_map[split]), 'r') as f:
            for l in f.readlines():
                ImageID = int(l)
                idx = (df[:, 1] == ImageID)
                if idx.sum() == 1:
                    jpeg_file = os.path.join(image_base, f"{ImageID}.jpg")
                    if f"{ImageID}.jpg" in avail_files:
                        r = {
                            "file_name": jpeg_file,
                            "image_id": ImageID,
                            "height": 1,
                            "width": 1,
                        }
                        r["label"] = df[idx, 2:12]
                        dicts.append(r)
    with open(cache_dir, 'wb') as f:
        torch.save(dicts, f)
    return dicts


def register_ava(name, dirname, split):
    DatasetCatalog.register(name, lambda: load_ava_instances(dirname, split))
    MetadataCatalog.get(name).set(
        thing_classes=CLASS_NAMES, dirname=dirname, split=split
    )
